package com.aloudmeta.grpc.support.metric;

import io.grpc.ForwardingServerCallListener;
import io.grpc.ServerCall.Listener;
import io.grpc.Status;
import io.micrometer.core.instrument.Counter;

import java.util.function.Consumer;
import java.util.function.Supplier;

/**
 * huaixin 2022/2/9 2:10 PM
 */
class MetricCollectingServerCallListener<Q> extends ForwardingServerCallListener.SimpleForwardingServerCallListener<Q> {

    private final Counter requestCounter;
    private final Supplier<Status.Code> responseCodeSupplier;
    private final Consumer<Status.Code> responseStatusTiming;

    /**
     * Creates a new delegating ServerCallListener that will wrap the given server call listener to collect metrics.
     *
     * @param delegate The original listener to wrap.
     * @param requestCounter The counter for incoming requests.
     * @param responseCodeSupplier The supplier of the response code.
     * @param responseStatusTiming The consumer used to time the processing duration along with a response status.
     */

    public MetricCollectingServerCallListener(
            final Listener<Q> delegate,
            final Counter requestCounter,
            final Supplier<Status.Code> responseCodeSupplier,
            final Consumer<Status.Code> responseStatusTiming) {

        super(delegate);
        this.requestCounter = requestCounter;
        this.responseCodeSupplier = responseCodeSupplier;
        this.responseStatusTiming = responseStatusTiming;
    }

    @Override
    public void onMessage(final Q requestMessage) {
        this.requestCounter.increment();
        super.onMessage(requestMessage);
    }

    @Override
    public void onComplete() {
        report(this.responseCodeSupplier.get());
        super.onComplete();
    }

    @Override
    public void onCancel() {
        report(Status.Code.CANCELLED);
        super.onCancel();
    }

    private void report(final Status.Code code) {
        this.responseStatusTiming.accept(code);
    }

}
